import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt


# 方程：θ''+b*θ'+c*sin(θ)=0
def pend(y, t, b, c):  # 令y为向量[ θ，ω ]
    theta, omega = y
    dydt = [omega, -b * omega - c * np.sin(theta)]  # 将θ’‘和θ‘表示出来
    return dydt  # 返回 dy/dt。 角度的导数就是角速度，角速度的导数是角加速度，也就是将其他几项移到等式右边。


b = 0.25
c = 5

y0 = [np.pi / 2, 0.0]  # 初始时假设单摆的初始角度为90度，角速度为0
t = np.linspace(0, 10, 101)
sol = odeint(pend, y0, t, args=(b, c))

plt.plot(t, sol[:, 0], 'b', label='theta(t)')
plt.plot(t, sol[:, 1], 'g', label='omega(t)')
plt.legend(loc='best')
plt.xlabel('t')
plt.grid()
plt.show()
